from fastapi import FastAPI, UploadFile, File, HTTPException
from fastapi.responses import JSONResponse
import pandas as pd
import numpy as np
from typing import Optional, Union, List
from pydantic import BaseModel
import io
import asyncio
from concurrent.futures import ProcessPoolExecutor
import modin.pandas as mpd  # 用于大数据处理的pandas替代品

app = FastAPI(title="Excel VLOOKUP Service")

# 全局进程池
process_pool = ProcessPoolExecutor()
# 全局缓存
df_cache = {}

class VlookupRequest(BaseModel):
    lookup_value: Union[str, int, float, List[Union[str, int, float]]]  # 支持单个值或值列表
    col_index_num: int
    range_lookup: Optional[bool] = False
    sheet_name: Optional[str] = None  # Excel sheet名称

def create_index_for_df(df: pd.DataFrame, lookup_column: str):
    """为DataFrame创建索引"""
    if not df.index.equals(df[lookup_column]):
        df.set_index(lookup_column, inplace=True)
        df.sort_index(inplace=True)

def batch_exact_lookup(df: pd.DataFrame, lookup_values, return_column: str):
    """批量精确查找"""
    return df.loc[lookup_values][return_column]

def batch_range_lookup(df: pd.DataFrame, lookup_values, return_column: str):
    """批量范围查找"""
    results = []
    for val in lookup_values:
        mask = df.index <= val
        if mask.any():
            idx = df.index[mask].max()
            results.append(df.loc[idx, return_column])
        else:
            results.append(None)
    return pd.Series(results, index=lookup_values)

@app.post("/vlookup/")
async def vlookup(
    params: VlookupRequest,
    lookup_file: UploadFile = File(...)
):
    try:
        # 获取文件内容
        file_content = await lookup_file.read()
        file_key = lookup_file.filename

        # 检查缓存
        if file_key not in df_cache:
            # 使用异步方式读取Excel
            loop = asyncio.get_event_loop()
            df = await loop.run_in_executor(
                process_pool,
                lambda: pd.read_excel(
                    io.BytesIO(file_content),
                    sheet_name=params.sheet_name,
                    engine='openpyxl'
                )
            )

            # 创建索引
            lookup_column = df.columns[0]
            create_index_for_df(df, lookup_column)

            # 存入缓存
            df_cache[file_key] = {
                'df': df,
                'timestamp': pd.Timestamp.now()
            }
        else:
            df = df_cache[file_key]['df']

        # 确保col_index_num有效
        if params.col_index_num < 1 or params.col_index_num > len(df.columns) + 1:
            raise HTTPException(
                status_code=400,
                detail=f"列索引 {params.col_index_num} 超出范围，表格共有 {len(df.columns)} 列"
            )

        return_column = df.columns[params.col_index_num - 2]  # -2是因为第一列已经是索引

        # 转换lookup_value为列表形式
        lookup_values = params.lookup_value if isinstance(params.lookup_value, list) else [params.lookup_value]

        try:
            if params.range_lookup:
                results = await loop.run_in_executor(
                    process_pool,
                    batch_range_lookup,
                    df, lookup_values, return_column
                )
            else:
                results = await loop.run_in_executor(
                    process_pool,
                    batch_exact_lookup,
                    df, lookup_values, return_column
                )

            # 清理超过1小时的缓存
            current_time = pd.Timestamp.now()
            df_cache.clear() if current_time - df_cache[file_key]['timestamp'] > pd.Timedelta(hours=1) else None

            return JSONResponse(
                content={
                    "status": "success",
                    "lookup_values": lookup_values,
                    "results": results.to_dict() if isinstance(results, pd.Series) else results
                }
            )

        except KeyError:
            raise HTTPException(
                status_code=404,
                detail=f"未找到匹配值: {lookup_values}"
            )

    except Exception as e:
        if isinstance(e, HTTPException):
            raise e
        raise HTTPException(
            status_code=500,
            detail=f"处理文件时发生错误: {str(e)}"
        )

@app.get("/health")
async def health_check():
    return {"status": "healthy"}

# 启动时创建进程池
@app.on_event("startup")
async def startup_event():
    global process_pool
    process_pool = ProcessPoolExecutor()

# 关闭时清理进程池
@app.on_event("shutdown")
async def shutdown_event():
    global process_pool
    process_pool.shutdown()
